use sea_orm::{Database, DatabaseConnection, DbErr, ConnectionTrait};
use std::env;
use tracing::{info, warn, error};
use url::Url;

pub async fn connect_database() -> Result<DatabaseConnection, DbErr> {
    // 从环境变量获取数据库URL，必须是PostgreSQL
    let database_url = env::var("DATABASE_URL")
        .expect("DATABASE_URL环境变量必须设置为PostgreSQL连接字符串");
    
    if !database_url.starts_with("postgres://") && !database_url.starts_with("postgresql://") {
        return Err(DbErr::Custom("DATABASE_URL必须是PostgreSQL连接字符串".to_string()));
    }
    
    info!("正在使用Sea-ORM连接PostgreSQL数据库: {}", database_url);
    
    // 尝试连接数据库
    match Database::connect(&database_url).await {
        Ok(db) => {
            info!("数据库连接成功");
            // 创建数据库表（如果不存在）
            create_tables_if_not_exists(&db).await?;
            Ok(db)
        }
        Err(e) => {
            // 检查是否是数据库不存在的错误
            let error_msg = e.to_string();
            if error_msg.contains("database") && error_msg.contains("does not exist") {
                warn!("数据库不存在，尝试自动创建...");
                
                // 尝试自动创建数据库
                match create_database_if_not_exists(&database_url).await {
                    Ok(_) => {
                        info!("数据库创建成功，重新尝试连接...");
                        // 重新连接数据库
                        let db = Database::connect(&database_url).await?;
                        // 创建数据库表（如果不存在）
                        create_tables_if_not_exists(&db).await?;
                        Ok(db)
                    }
                    Err(create_err) => {
                        error!("数据库自动创建失败: {}", create_err);
                        Err(create_err)
                    }
                }
            } else {
                // 其他连接错误，直接返回
                error!("数据库连接失败: {}", e);
                Err(e)
            }
        }
    }
}

async fn create_database_if_not_exists(database_url: &str) -> Result<(), DbErr> {
    // 解析数据库URL
    let url = Url::parse(database_url)
        .map_err(|e| DbErr::Custom(format!("解析数据库URL失败: {}", e)))?;
    
    // 获取数据库名称
    let db_name = url.path().trim_start_matches('/');
    if db_name.is_empty() {
        return Err(DbErr::Custom("数据库URL中未指定数据库名称".to_string()));
    }
    
    // 构建连接到postgres默认数据库的URL（用于创建新数据库）
    let mut postgres_url = url.clone();
    postgres_url.set_path("/postgres");
    let postgres_url_str = postgres_url.to_string();
    
    info!("连接到postgres默认数据库以创建目标数据库: {}", db_name);
    
    // 连接到postgres默认数据库
    let postgres_db = Database::connect(&postgres_url_str).await
        .map_err(|e| DbErr::Custom(format!("连接到postgres默认数据库失败: {}", e)))?;
    
    // 检查数据库是否已存在
    use sea_orm::Statement;
    let check_db_sql = format!(
        "SELECT 1 FROM pg_database WHERE datname = '{}'",
        db_name
    );
    
    let statement = Statement::from_string(sea_orm::DatabaseBackend::Postgres, check_db_sql);
    let result = postgres_db.query_one(statement).await;
    
    match result {
        Ok(Some(_)) => {
            info!("数据库 '{}' 已存在，无需创建", db_name);
            Ok(())
        }
        Ok(None) => {
            // 数据库不存在，创建它
            info!("创建数据库: {}", db_name);
            let create_db_sql = format!("CREATE DATABASE \"{}\"", db_name);
            let create_statement = Statement::from_string(sea_orm::DatabaseBackend::Postgres, create_db_sql);
            
            postgres_db.execute(create_statement).await
                .map_err(|e| DbErr::Custom(format!("创建数据库失败: {}", e)))?;
            
            info!("数据库 '{}' 创建成功", db_name);
            Ok(())
        }
        Err(e) => {
            Err(DbErr::Custom(format!("检查数据库是否存在时出错: {}", e)))
        }
    }
}

async fn create_tables_if_not_exists(db: &DatabaseConnection) -> Result<(), DbErr> {
    use sea_orm::Statement;
    
    info!("检查并创建数据库表...");
    
    // 创建files表
    let create_files_table = Statement::from_string(
        sea_orm::DatabaseBackend::Postgres,
        r#"
        CREATE TABLE IF NOT EXISTS files (
            id SERIAL PRIMARY KEY,
            file_path TEXT UNIQUE NOT NULL,
            file_name TEXT NOT NULL,
            file_size BIGINT NOT NULL,
            file_hash TEXT NOT NULL,
            field_order JSONB,
            created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
            updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
        )
        "#.to_string()
    );
    
    db.execute(create_files_table).await?;
    info!("files表检查完成");
    
    // 创建excel_data表
    let create_excel_data_table = Statement::from_string(
        sea_orm::DatabaseBackend::Postgres,
        r#"
        CREATE TABLE IF NOT EXISTS excel_data (
            id SERIAL PRIMARY KEY,
            file_id INTEGER NOT NULL REFERENCES files(id) ON DELETE CASCADE,
            import_time TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
            row_number INTEGER NOT NULL,
            sheet_name TEXT NOT NULL DEFAULT 'Sheet1',
            data_json JSONB NOT NULL,
            search_text TEXT NOT NULL
        )
        "#.to_string()
    );
    
    db.execute(create_excel_data_table).await?;
    info!("excel_data表检查完成");
    
    // 创建索引
    let indexes = vec![
        "CREATE INDEX IF NOT EXISTS idx_excel_data_search_text ON excel_data(search_text)",
        "CREATE INDEX IF NOT EXISTS idx_excel_data_file_id ON excel_data(file_id)",
        "CREATE INDEX IF NOT EXISTS idx_excel_data_import_time ON excel_data(import_time)",
        "CREATE INDEX IF NOT EXISTS idx_excel_data_data_json ON excel_data USING GIN (data_json)",
        "CREATE INDEX IF NOT EXISTS idx_files_file_path ON files(file_path)",
        "CREATE INDEX IF NOT EXISTS idx_files_file_hash ON files(file_hash)",
    ];
    
    for index_sql in indexes {
        let statement = Statement::from_string(sea_orm::DatabaseBackend::Postgres, index_sql.to_string());
        db.execute(statement).await?;
    }
    
    info!("数据库索引检查完成");
    info!("数据库表和索引初始化完成");
    
    Ok(())
}